import os
import sys
from src.core import log
from alembic import command
from alembic.config import Config
from ..library.manage import db_manager, Base

class AlembicManager:
    def __init__(self):
        # 设置 Alembic 配置路径
        self.alembic_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'alembic')
        self.config_file = os.path.join(self.alembic_dir, 'alembic.ini')
        
        # 确保目录存在
        os.makedirs(self.alembic_dir, exist_ok=True)
        os.makedirs(os.path.join(self.alembic_dir, 'versions'), exist_ok=True)
        
        # 初始化配置
        self.config = Config(self.config_file)
        self.config.set_main_option('script_location', self.alembic_dir)
        
        # 设置数据库URL（使用第一个数据库）
        db_names = list(db_manager.engines.keys())
        if db_names:
            main_engine = db_manager.engines[db_names[0]]
            self.config.set_main_option('sqlalchemy.url', str(main_engine.url))
        
        # 创建必要的文件
        self._create_alembic_files()

    def _create_alembic_files(self):
        """创建必要的 Alembic 文件"""
        # 创建 env.py
        env_py_path = os.path.join(self.alembic_dir, 'env.py')
        if not os.path.exists(env_py_path):
            with open(env_py_path, 'w', encoding='utf-8') as f:
                f.write('''# coding: utf-8
import os
import sys
from logging.config import fileConfig
from alembic import context
from sqlalchemy import engine_from_config, pool

# 添加项目根目录到 Python 路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

# 导入您的数据库管理器和模型
from src.core.database import db_manager, Base

# 设置目标元数据
target_metadata = Base.metadata

def run_migrations_offline():
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
        render_as_batch=True,
        compare_type=True,
        compare_server_default=True
    )
    with context.begin_transaction():
        context.run_migrations()

def run_migrations_online():
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )
    
    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            render_as_batch=True,
            compare_type=True,
            compare_server_default=True
        )
        with context.begin_transaction():
            context.run_migrations()

if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()
''')

        # 创建 script.py.mako
        script_mako_path = os.path.join(self.alembic_dir, 'script.py.mako')
        if not os.path.exists(script_mako_path):
            with open(script_mako_path, 'w', encoding='utf-8') as f:
                f.write('''"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from alembic import op
import sqlalchemy as sa
${imports if imports else ""}

# revision identifiers, used by Alembic.
revision = ${repr(up_revision)}
down_revision = ${repr(down_revision)}
branch_labels = ${repr(branch_labels)}
depends_on = ${repr(depends_on)}


def upgrade():
    ${upgrades if upgrades else "pass"}


def downgrade():
    ${downgrades if downgrades else "pass"}
''')

    def init_alembic(self):
        """初始化 Alembic"""
        try:
            # 检查是否已初始化
            if not os.path.exists(os.path.join(self.alembic_dir, 'env.py')):
                log.INFO("初始化 Alembic 迁移系统...")
                # 创建基础文件
                self._create_alembic_files()
                log.INFO("Alembic 初始化完成")
            return True
        except Exception as e:
            log.ERROR(f"Alembic 初始化失败: {e}")
            return False

    def generate_migration(self, message="auto migration"):
        """生成迁移脚本"""
        try:
            log.INFO(f"生成迁移脚本: {message}")
            command.revision(self.config, autogenerate=True, message=message)
            log.INFO("迁移脚本生成成功")
            return True
        except Exception as e:
            log.ERROR(f"生成迁移脚本失败: {e}")
            return False

    def apply_migrations(self, revision="head"):
        """应用迁移到所有数据库"""
        try:
            log.INFO(f"应用迁移到数据库: {revision}")
            
            # 应用到所有数据库
            success_count = 0
            for db_name, engine in db_manager.engines.items():
                try:
                    log.INFO(f"正在应用到数据库: {db_name}")
                    
                    # 临时修改配置使用当前数据库
                    temp_config = Config(self.config_file)
                    temp_config.set_main_option('script_location', self.alembic_dir)
                    temp_config.set_main_option('sqlalchemy.url', str(engine.url))
                    
                    # 应用迁移
                    command.upgrade(temp_config, revision)
                    success_count += 1
                    log.INFO(f"数据库 {db_name} 迁移成功")
                    
                except Exception as e:
                    log.ERROR(f"数据库 {db_name} 迁移失败: {e}")
            
            log.INFO(f"迁移完成，成功: {success_count}/{len(db_manager.engines)}")
            return success_count > 0
            
        except Exception as e:
            log.ERROR(f"应用迁移失败: {e}")
            return False

    def get_current_revision(self):
        """获取当前迁移版本"""
        try:
            from alembic.migration import MigrationContext
            import sqlalchemy as sa
            
            # 检查第一个数据库的版本
            if db_manager.engines:
                first_engine = list(db_manager.engines.values())[0]
                with first_engine.connect() as conn:
                    context = MigrationContext.configure(conn)
                    return context.get_current_revision()
            return None
        except:
            return None

    def get_migration_history(self):
        """获取迁移历史"""
        try:
            from alembic.script import ScriptDirectory
            script = ScriptDirectory.from_config(self.config)
            return [(rev.revision, rev.doc) for rev in script.walk_revisions()]
        except Exception as e:
            log.ERROR(f"获取迁移历史失败: {e}")
            return []

    def auto_migrate(self):
        """自动检测并执行迁移"""
        try:
            # 检查是否有未应用的迁移
            current_rev = self.get_current_revision()
            if current_rev is None:
                log.INFO("首次迁移，创建初始版本")
                self.generate_migration("initial migration")
            
            # 生成新的迁移（如果有变化）
            self.generate_migration("auto update")
            
            # 应用迁移
            return self.apply_migrations()
            
        except Exception as e:
            log.ERROR(f"自动迁移失败: {e}")
            return False

# 全局实例
alembic_manager = AlembicManager()